package com.sunday.common.mq.kafka.brave.core;

import brave.kafka.clients.KafkaTracing;
import io.micrometer.observation.Observation;
import org.apache.kafka.clients.producer.Producer;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.boot.autoconfigure.kafka.KafkaConnectionDetails;
import org.springframework.kafka.core.DefaultKafkaProducerFactory;
import org.springframework.lang.Nullable;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Kafka Producer 埋点
 * 执行过程:
 * {@link com.sunday.common.mq.kafka.study.spring.e20_Brave_Trace.KafkaBraveConfiguration#kafkaProducerFactory(KafkaConnectionDetails, ObjectProvider, KafkaTracing) }
 * {@link org.springframework.kafka.core.KafkaTemplate#observeSend(ProducerRecord)}
 * {@link org.springframework.kafka.core.KafkaTemplate#doSend(ProducerRecord, Observation)}
 * {@link org.springframework.kafka.core.KafkaTemplate#getTheProducer(String)}
 * {@link org.springframework.kafka.core.DefaultKafkaProducerFactory#createProducer()}
 * {@link org.springframework.kafka.core.DefaultKafkaProducerFactory#doCreateProducer(String)}
 *
 * @author zsy
 * @see brave.kafka.clients.TracingProducer 参考
 * @since 2023/11/17
 */
public class TracingKafkaProducerFactory<K, V> extends DefaultKafkaProducerFactory<K, V> {

    private final KafkaTracing kafkaTracing;

    /**
     * 首先因为 org.springframework.kafka.core.KafkaTemplate#getTheProducer(java.lang.String)
     * 方法中每次获取 producer 都会涉及到 createProducer 和 createNonTransactionalProducer 方法，导致了每次发送消息都会来调用方法
     * 也产生了每次都要构建 kafkaTracing.producer，所以采用进行所属的 producer 对应绑定里链路后的 trace producer的缓存
     * PS ： 这里 consumer不存在这个问题，所以没有启用缓存
     */
    private volatile Map<Producer, Producer> cache = new ConcurrentHashMap<>();

    public TracingKafkaProducerFactory(Map<String, Object> configs, KafkaTracing kafkaTracing) {
        super(configs);
        this.kafkaTracing = kafkaTracing;
    }

    @Override
    public Producer<K, V> createProducer(@Nullable String txIdPrefixArg) {
        Producer<K, V> producer = super.createProducer(txIdPrefixArg);
        return cache.computeIfAbsent(producer, p -> kafkaTracing.producer(p));
    }

    @Override
    public Producer<K, V> createNonTransactionalProducer() {
        Producer<K, V> producer = super.createNonTransactionalProducer();
        return cache.computeIfAbsent(producer, p -> kafkaTracing.producer(p));
    }


}
